import aiohttp
import json
import asyncio
import logging
from typing import Dict, Any, Optional

from .ai_base import BaseAIService, GenerationRequest, GenerationResponse, AIServiceProvider, AIServiceType
from app.core.error_handler import robust_executor

logger = logging.getLogger(__name__)

class JimengImageFusionService:
    """即梦AI图片融合专用服务"""
    
    def __init__(self, jimeng_service):
        self.jimeng_service = jimeng_service
        self.image_fusion_req_key = "jimeng_i2i_v30"
    
    async def generate_image_fusion(self, request: GenerationRequest) -> GenerationResponse:
        """图片融合生成 - 带错误处理和重试"""
        return await robust_executor.execute_with_retry(
            self._generate_image_fusion_internal,
            "jimeng_ai_fusion",
            request
        )
    
    async def _generate_image_fusion_internal(self, request: GenerationRequest) -> GenerationResponse:
        """内部图片融合方法 - 使用异步API"""
        try:
            # 构建请求参数
            payload = self._build_image_fusion_payload(request)
            query_params = {
                'Action': 'CVSync2AsyncSubmitTask',
                'Version': '2022-08-31'
            }
            
            # 生成签名头
            body_json = json.dumps(payload)
            headers = self.jimeng_service.signature_helper.generate_v4_signature('POST', query_params, body_json)
            
            # 发送请求
            query_string = self.jimeng_service.signature_helper._format_query(query_params)
            request_url = f"{self.jimeng_service.endpoint}?{query_string}"
            
            async with aiohttp.ClientSession() as session:
                async with session.post(request_url, headers=headers, data=body_json) as response:
                    response_text = await response.text()
                    response_data = json.loads(response_text)
                    
                    if response.status == 200 and response_data.get('code') == 10000:
                        # 异步API，获取任务ID
                        data = response_data.get('data', {})
                        task_id = data.get('task_id')
                        
                        if not task_id:
                            raise Exception("图片融合未返回任务ID")
                        
                        # 轮询获取结果，与视频生成流程一致
                        fusion_result = await self.jimeng_service._poll_task_result(task_id, self.image_fusion_req_key)
                        
                        if fusion_result and fusion_result.get("status") == "done":
                            image_urls = fusion_result.get("image_urls")
                            base64_data = fusion_result.get("binary_data_base64")

                            if image_urls:
                                final_url = image_urls[0]
                            elif base64_data and base64_data[0]:
                                final_url = f"data:image/jpeg;base64,{base64_data[0]}"
                            else:
                                raise Exception("图片融合任务完成但未返回结果图片(URL或Base64)")
                                
                            return GenerationResponse(
                                success=True,
                                result_url=final_url,
                                cost=0.02,  # 图片融合基础成本
                                metadata={
                                    "task_id": task_id,
                                    "fusion_type": "image_fusion",
                                    "api_type": "async",
                                    "result_count": 1,
                                    "raw_result": fusion_result
                                }
                            )
                        else:
                            return GenerationResponse(
                                success=False, 
                                error_message="图片融合任务失败或超时"
                            )
                    else:
                        error_msg = response_data.get('message', f'HTTP {response.status}')
                        return GenerationResponse(success=False, error_message=f"图片融合请求失败: {error_msg}")
                        
        except Exception as e:
            return GenerationResponse(success=False, error_message=f"图片融合异常: {str(e)}")
    
    def _build_image_fusion_payload(self, request: GenerationRequest) -> Dict[str, Any]:
        """构建图片融合请求payload"""
        payload = {
            "req_key": self.image_fusion_req_key,
            "prompt": request.prompt,
            "scale": request.params.get("fusion_strength", 0.6),
            "seed": request.params.get("seed", -1)
        }
        
        # 根据用户传入的比例参数映射到对应的宽高尺寸
        aspect_ratio = request.params.get("aspect_ratio", "16:9")
        width, height = self._get_dimensions_by_aspect_ratio(aspect_ratio)
        payload["width"] = width
        payload["height"] = height
        
        # 处理图片输入 - 根据即梦AI文档，只支持单张图片
        template_url = request.params.get("template_image_url")
        template_base64 = request.params.get("template_image_base64")
        user_url = request.params.get("user_image_url")
        user_base64 = request.params.get("user_image_base64")
        
        # 选择策略：优先模板图，如果没有则用用户图
        if template_url:
            payload["image_urls"] = [template_url]
        elif user_url:
            payload["image_urls"] = [user_url] 
        elif template_base64:
            payload["binary_data_base64"] = [template_base64]
        elif user_base64:
            payload["binary_data_base64"] = [user_base64]
        else:
            # 兼容旧参数名
            user_urls = request.params.get("user_image_urls") or request.params.get("image_urls")
            legacy_base64 = request.params.get("binary_data_base64")
            
            if user_urls:
                payload["image_urls"] = [user_urls[0]]  # 只使用第一张
            elif legacy_base64:
                payload["binary_data_base64"] = [legacy_base64[0]]  # 只使用第一张
        
        return payload
    
    def _get_dimensions_by_aspect_ratio(self, aspect_ratio: str) -> tuple[int, int]:
        """根据宽高比获取对应的宽高尺寸"""
        aspect_ratio_mapping = {
            "1:1": (1328, 1328),    # 1:1
            "4:3": (1472, 1104),    # 4:3
            "3:2": (1584, 1056),    # 3:2
            "16:9": (1664, 936),    # 16:9
            "21:9": (2016, 864),    # 21:9
        }
        return aspect_ratio_mapping.get(aspect_ratio, (1664, 936))  # 默认16:9